package com.zeng.user.utils;

import java.util.*;

import static com.zeng.user.utils.ConsistentHash.SERVERS;

/**
 * @author zyb
 * @date 2021-05-17 21:16
 */
public class Consistent1Hash {

    private final Set<String> sourceServers;
    private final int virtualSize;

    private final SortedMap<Integer, String> cachedMap = new TreeMap<>();
    private final Map<String, String> virtualMap = new TreeMap<>();


    public Consistent1Hash(Set<String> servers, int virtualSize) {
        this.virtualSize = virtualSize;
        this.sourceServers = new HashSet<>(servers);
        for (String server : sourceServers) {
            addServer(server);
        }
    }

    public void addServer(String server) {
        for (int i = 0; i < virtualSize; i++) {
            String virtualKey = server + "#" + i;
            virtualMap.put(virtualKey, server);
            int hash = hashCode(virtualKey);
            cachedMap.put(hash, virtualKey);
        }
    }

    public void removeServer(String server) {
        List<String> virtualKeys = new LinkedList<>();
        virtualMap.forEach((key, actualServer) -> {
            if (server.equals(actualServer)) {
                virtualKeys.add(key);
            }
        });
        for (String virtualKey : virtualKeys) {
            int hash = hashCode(virtualKey);
            cachedMap.remove(hash);
            virtualMap.remove(virtualKey);

        }
    }


    public String getServer(String key) {
        int hash = hashCode(key);
        SortedMap<Integer, String> tailMap = cachedMap.tailMap(hash);
        if (tailMap.isEmpty()) {
            tailMap = cachedMap;
        }
        String virtualKey = cachedMap.get(tailMap.firstKey());
        return virtualMap.get(virtualKey);
    }

    private static int hashCode(String source) {
        return Math.abs(source.hashCode());

    }

    public static void main(String[] args) {
        Consistent1Hash consistentHash = new Consistent1Hash(SERVERS, 10);
        String server = consistentHash.getServer("test");
        System.out.println(server);
    }
}
